import glob
import os
import shutil
import random


def split_data(file_path, train_dir, valid_dir, train_rate, val_rate, test_rate):
    """
    划分数据集合
    :param file_path:
    :param train_dir:
    :param valid_dir:
    :param train_rate:
    :param val_rate:
    :param test_rate:
    :return:
    """
    # 计算文件有关的信息
    each_label = []
    image_path = os.path.join(file_path, "images")
    label_path = os.path.join(file_path, "labels")
    each_image = glob.glob(os.path.join(image_path, '*.jpg'))
    for image_path in each_image:
        base_name = os.path.basename(image_path).split('.')[0]
        each_label.append(label_path + '/' + base_name + '.txt')
    data = list(zip(each_image, each_label))
    total = len(each_image)

    # 随机打乱数据
    random.shuffle(data)
    each_class_image, each_class_label = zip(*data)
    train_images = each_class_image[0:int(train_rate * total)]
    val_images = each_class_image[int(train_rate * total):int((train_rate + val_rate) * total)]
    test_images = each_class_image[int((train_rate + val_rate) * total):]
    train_labels = each_class_label[0:int(train_rate * total)]
    val_labels = each_class_label[int(train_rate * total):int((train_rate + val_rate) * total)]
    test_labels = each_class_label[int((train_rate + val_rate) * total):]

    os.makedirs(os.path.join(train_dir, 'images'), exist_ok=True)
    os.makedirs(os.path.join(train_dir, 'labels'), exist_ok=True)
    os.makedirs(os.path.join(valid_dir, 'images'), exist_ok=True)
    os.makedirs(os.path.join(valid_dir, 'labels'), exist_ok=True)

    for i in range(len(train_images)):
        image_basename = os.path.basename(train_images[i])
        label_basename = os.path.basename(train_labels[i])
        shutil.copy(train_images[i], os.path.join(train_dir, 'images', image_basename))
        shutil.copy(train_labels[i], os.path.join(train_dir, 'labels', label_basename))

    for j in range(len(val_images)):
        image_basename = os.path.basename(val_images[j])
        label_basename = os.path.basename(val_labels[j])
        shutil.copy(val_images[j], os.path.join(valid_dir, 'images', image_basename))
        shutil.copy(val_labels[j], os.path.join(valid_dir, 'labels', label_basename))


if __name__ == '__main__':
    split_data("D:/project/dataset/instrument/all", "D:/project/dataset/instrument/train",
               "D:/project/dataset/instrument/valid", 0.8, 0.1, 0.1)
